import os
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
from contextlib import contextmanager
import pymongo
import redis
from urllib.parse import quote_plus
import logging

# 日志配置
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库配置
MYSQL_HOST = os.getenv("MYSQL_HOST", "localhost")
MYSQL_PORT = os.getenv("MYSQL_PORT", "3306")   # 默认MySQL端口
MYSQL_USER = os.getenv("MYSQL_USER", "root")   # 数据库用户名
MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "password123")  # 数据库密码
MYSQL_DB = os.getenv("MYSQL_DB", "spider_data")  # 数据库名称

MONGODB_HOST = os.getenv("MONGODB_HOST", "localhost")
MONGODB_PORT = os.getenv("MONGODB_PORT", "27017")    # 默认MongoDB端口
MONGODB_USER = os.getenv("MONGODB_USER", "admin")    # MongoDB用户名
MONGODB_PASSWORD = os.getenv("MONGODB_PASSWORD", "password123")  # MongoDB密码
MONGODB_DB = os.getenv("MONGODB_DB", "spider_data")  # MongoDB数据库名称

REDIS_HOST = os.getenv("REDIS_HOST", "redis")
REDIS_PORT = os.getenv("REDIS_PORT", "6379")
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "")
REDIS_DB = os.getenv("REDIS_DB", "0")

# MySQL 连接
mysql_uri = f"mysql+pymysql://{MYSQL_USER}:{quote_plus(MYSQL_PASSWORD)}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DB}?charset=utf8mb4"
engine = create_engine(mysql_uri, pool_size=5, max_overflow=10, pool_recycle=3600, echo=False)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
session_factory = scoped_session(SessionLocal)

Base = declarative_base()

@contextmanager
def get_db_session():
    """数据库会话上下文管理器"""
    session = session_factory()
    try:
        yield session
        session.commit()
    except Exception as e:
        session.rollback()
        logger.error(f"数据库操作异常: {str(e)}")
        raise
    finally:
        session.close()

# MongoDB 连接
def get_mongodb_client():
    """获取MongoDB客户端"""
    try:
        mongo_uri = f"mongodb://{MONGODB_USER}:{quote_plus(MONGODB_PASSWORD)}@{MONGODB_HOST}:{MONGODB_PORT}/{MONGODB_DB}"
        client = pymongo.MongoClient(mongo_uri)
        return client
    except Exception as e:
        logger.error(f"MongoDB连接异常: {str(e)}")
        raise

def get_mongodb_db():
    """获取MongoDB数据库"""
    client = get_mongodb_client()
    return client[MONGODB_DB]

# Redis 连接
def get_redis_client():
    """获取Redis客户端"""
    try:
        pool = redis.ConnectionPool(
            host=REDIS_HOST,
            port=int(REDIS_PORT),
            password=REDIS_PASSWORD if REDIS_PASSWORD else None,
            db=int(REDIS_DB),
            decode_responses=True
        )
        return redis.Redis(connection_pool=pool)
    except Exception as e:
        logger.error(f"Redis连接异常: {str(e)}")
        raise

def init_db():
    """初始化数据库"""
    from src.database.models import Base
    Base.metadata.create_all(bind=engine)
    logger.info("数据库表初始化完成") 